#! /usr/bin/python2.5
"""XXX TODO:
    * Launch a thread to download the pages so GUI stays responsive.
        + wxPython has something built-in for this.
    * Have a progress gauge and a running count of the number of files
      downloaded.
        + Use ProgressDialog and its Pulse() method.
        + Probably want a subclass that acts as a controller.
        + Can use subclass to store final results.
        + Do support halting of the downloading.
    * Make downloading multi-threaded.
    * Add an audible bell when downloads finish.
    * Let wxPython redirect output.
    * Have a retry button that tries downloading again just past the last
      failed attempt.

"""
from __future__ import with_statement

import wx
from wx.lib.stattext import GenStaticText

from contextlib import contextmanager
from itertools import count
import os
import os.path
import sys
import urllib2
from urllib import unquote, quote
from urlparse import urlsplit


class HaltedDownloadException(Exception):

    def __init__(self, number, title, message):
        self.number = number
        self.title = title
        self.message = message
        Exception.__init__(self, title + ': ' + message)


@contextmanager
def disable_buttons(*buttons):
    try:
        for button in buttons:
            button.Enable(False)
        yield
    finally:
        for button in buttons:
            button.Enable(True)


class App(wx.App):

	def OnInit(self):
		"""Set up application."""
		self.frame = MainWindow()
		self.frame.Show(True)
		self.SetTopWindow(self.frame)
		return True

class MainWindow(wx.Frame):
	"""Create a window that takes in a URL (with a selection) and a download
	button."""

	def __init__(self):
		wx.Frame.__init__(self, None, wx.ID_ANY, "Sequential Downloader")

		panel = wx.Panel(self)
                # File URL.
		url_text = wx.StaticText(panel, wx.ID_ANY, "URL (w/ selection)")
		self.url_ctrl = wx.TextCtrl(panel, wx.ID_ANY, size=(500,-1))
                # Directory.
                dir_button = wx.Button(panel, wx.ID_ANY, "Dir ...")
                self.dir_text = GenStaticText(panel, wx.ID_ANY, os.getcwd())
                self.Bind(wx.EVT_BUTTON, self.change_dir, dir_button)
                # Download.
		download_button = wx.Button(panel, wx.ID_ANY, "Download")
		self.Bind(wx.EVT_BUTTON, self.download, download_button)
		download_button.SetDefault()

                self.buttons = [dir_button, download_button]

		input_sizer = wx.FlexGridSizer(2, 2, 5, 10)
		input_sizer.AddGrowableCol(0)
		input_sizer.AddGrowableCol(1)
		input_sizer.Add(url_text, flag=wx.ALIGN_RIGHT)
		input_sizer.Add(self.url_ctrl, flag=wx.EXPAND)
                input_sizer.Add(dir_button, flag=wx.ALIGN_CENTER)
                input_sizer.Add(self.dir_text, flag=wx.EXPAND)

                button_sizer = wx.FlexGridSizer(1, 1)
                button_sizer.Add(download_button, flag=wx.ALIGN_CENTER)

		frame_sizer = wx.GridBagSizer(10, 10)
		frame_sizer.Add(input_sizer, (0,0))
		frame_sizer.Add(button_sizer, (1, 0), flag=wx.ALIGN_CENTER)
		panel.SetSizerAndFit(frame_sizer)
		self.SetClientSize(panel.GetSize())

        def change_dir(self, event=None):
            """Change the working directory."""
            dir_dialog = wx.DirDialog(self, "Choose a directory to save to",
                                        os.getcwd())
            dir_dialog.ShowModal()
            directory = dir_dialog.GetPath()
            os.chdir(directory)
            self.dir_text.SetLabel(directory)

	def download(self, event=None):
           """Download a sequence of files based on the URL (and its selection)
           from the self.url_ctrl text field."""
           with disable_buttons(*self.buttons):
               # Get input.
               start_url = unquote(self.url_ctrl.GetValue())
               varying_part = self.url_ctrl.GetStringSelection()
               if not varying_part:
                   need_selection_dialog = wx.MessageDialog(self,
                                       "The changing number was not selected.",
                                       caption="Selection Needed",
                                       style=wx.OK | wx.ICON_ERROR)
                   need_selection_dialog.ShowModal()
                   return
               # Figure out URL template.
               root_url, base_name = start_url.rsplit('/', 1)
               partition = base_name.rpartition(varying_part)
               base_name = partition[0] + '%s' + partition[2]
               if base_name.endswith('.html'):
                   base_name = base_name.replace('.html', '.jpg')
               start = int(varying_part)
               try:
                   result = sequential_download(root_url, base_name, start,
                                                   zfill=len(varying_part))
               except HaltedDownloadException, exc:
                   error_dialog = wx.MessageDialog(self, str(exc.args[0]),
                           caption=str(exc.title), style=wx.OK | wx.ICON_ERROR)
                   error_dialog.ShowModal()
                   return
               else:
                   msg = """Downloading is complete.

                   Largest number downloaded: %s
                   Large number attempted: %s""" % result
                   done_dialog = wx.MessageDialog(self, msg,
                                               caption="Done Downloading",
                                               style=wx.OK | wx.ICON_INFORMATION)
                   done_dialog.ShowModal()
                   return


def sequential_download(root_url, base_name, start, zfill):
    """Download URLs starting at 'start'."""
    previous_size = 0
    downloaded_num = start - 1
    skipped_num = downloaded_num
    for num in count(start):
        file_name = base_name % str(num).zfill(zfill)
        url = '/'.join([root_url, file_name])
        # XXX Keep spinner going.
        print 'Downloading %s ...' % url,
        sys.stdout.flush()
        parsed_url = urlsplit(url)
        url = "%s://%s%s" % (parsed_url.scheme, parsed_url.netloc,
                            quote(parsed_url.path))
        try:
            req = urllib2.Request(url)
            parsed_url = urlsplit(url)
            referrer = '%s://%s' % (parsed_url.scheme, parsed_url.netloc)
            req.add_header('Referer', referrer)
            downloaded = urllib2.urlopen(req)
        except urllib2.URLError, exc:
            skipped_num = num if num > skipped_num else skipped_num
            # XXX Update largest number that didn't get downloaded.
            print "skipped (%s)" % exc
            if skipped_num > downloaded_num + 4:
                return downloaded_num, skipped_num
            else:
                continue
        else:
            downloaded_num = num if num > downloaded_num else downloaded_num
            data = downloaded.read()
            data_amount = len(data)
            if data_amount == previous_size:
                raise HaltedDownloadException(num, "File Same Size",
                        "The file %s is the same size as the previous file." %
                            file_name)
            previous_size = data_amount
            if os.path.exists(file_name):
                raise HaltedDownloadException(num, "File Exists",
                                    "The file %s already exists." % file_name)
            with open(file_name, 'wb') as file_:
                file_.write(data)
            print "saved"


def main():
    app = App(redirect=False)
    app.MainLoop()



if __name__ == '__main__':
    main()
